package com.red.utils;

import org.springframework.core.io.ClassPathResource;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.*;

@Component
public class SensitiveWordFilter {

    private static final String ENCODING = "UTF-8";
    private static final String REPLACEMENT = "***";
    private TrieNode rootNode = new TrieNode();

    // Trie树节点
    private class TrieNode {
        private boolean end;
        private Map<Character, TrieNode> subNodes = new HashMap<>();

        void addSubNode(Character key, TrieNode node) {
            subNodes.put(key, node);
        }

        TrieNode getSubNode(Character key) {
            return subNodes.get(key);
        }

        boolean isEnd() {
            return end;
        }

        void setEnd(boolean end) {
            this.end = end;
        }
    }

    @PostConstruct
    public void init() throws IOException {
        // 从classpath加载敏感词文件
        ClassPathResource resource = new ClassPathResource("sensitive-words.txt");
        InputStream inputStream = resource.getInputStream();
        BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream, ENCODING));

        String keyword;
        while ((keyword = reader.readLine()) != null) {
            addKeyword(keyword.trim());
        }
        reader.close();
    }

    private void addKeyword(String keyword) {
        TrieNode tempNode = rootNode;
        for (int i = 0; i < keyword.length(); i++) {
            char c = keyword.charAt(i);
            TrieNode node = tempNode.getSubNode(c);
            if (node == null) {
                node = new TrieNode();
                tempNode.addSubNode(c, node);
            }
            tempNode = node;
            if (i == keyword.length() - 1) {
                tempNode.setEnd(true);
            }
        }
    }

    public String filter(String text) {
        if (text == null || text.trim().isEmpty()) {
            return text;
        }

        StringBuilder result = new StringBuilder();
        TrieNode tempNode = rootNode;
        int begin = 0;
        int position = 0;

        while (position < text.length()) {
            char c = text.charAt(position);
            tempNode = tempNode.getSubNode(c);
            if (tempNode == null) {
                result.append(text.charAt(begin));
                position = begin + 1;
                begin = position;
                tempNode = rootNode;
            } else if (tempNode.isEnd()) {
                result.append(REPLACEMENT);
                position = position + 1;
                begin = position;
                tempNode = rootNode;
            } else {
                position++;
            }
        }
        result.append(text.substring(begin));
        return result.toString();
    }
}